Context Patterns

// Read-only instruction generation - use ReadonlyContext
const safeInstruction = (ctx: ReadonlyContext): string => {
  return `Hello ${ctx.state.user?.name || "User"}! Session: ${ctx.invocationId}`;
};

// State modification in callbacks - use CallbackContext
async function updateUserState(ctx: CallbackContext) {
  ctx.state.lastActive = new Date().toISOString();
  await ctx.saveArtifact("activity.json", { text: "User active" });
}

// Tool with memory access - use ToolContext
async function smartTool(params: any, ctx: ToolContext) {
  const memories = await ctx.searchMemory(params.query);
  ctx.state.lastSearch = memories;
  return { found: memories.memories?.length || 0 };
}

// Full agent implementation - use InvocationContext
export class SecureAgent extends BaseAgent {
  protected async *runAsyncImpl(ctx: InvocationContext) {
    // Full framework access with all services
    yield* this.processSecurely(ctx);
  }
}

// Good - ReadonlyContext for safe read operations
function analyzeState(ctx: ReadonlyContext) {
  return {
    userCount: Object.keys(ctx.state).filter(k => k.startsWith('user.')).length,
    sessionAge: Date.now() - (ctx.state.sessionStart || Date.now())
  };
}

// Avoid - ToolContext when you only need to read state
function badAnalyzeState(ctx: ToolContext) {
  // Unnecessarily powerful context for read-only operation
  return { userCount: 1 };
}

function manageStateScopes(ctx: CallbackContext) {
  // Session-specific (cleared when session ends)
  ctx.state.currentTask = "processing";
  ctx.state.sessionProgress = 0.5;

  // User-specific (persists across sessions)
  ctx.state["user.preferences"] = { theme: "dark", language: "en" };
  ctx.state["user.profile"] = { name: "John", level: "expert" };

  // App-specific (shared across all users)
  ctx.state["app.version"] = "1.2.0";
  ctx.state["app.features"] = ["feature1", "feature2"];

  // Temporary (explicitly cleaned up)
  ctx.state["temp.processing"] = { status: "active", startTime: Date.now() };
  ctx.state["temp.uploads"] = ["file1.txt", "file2.pdf"];
}

function validateAndUpdateState(ctx: CallbackContext, updates: any) {
  // Validate state schema
  const schema = {
    userPreferences: { theme: "string", language: "string" },
    currentTask: "string",
    progress: "number"
  };

  // Backup current state
  const backup = JSON.stringify(ctx.state);

  try {
    // Validate updates
    for (const [key, value] of Object.entries(updates)) {
      if (schema[key]) {
        const expectedType = schema[key];
        if (typeof value !== expectedType) {
          throw new Error(`Invalid type for ${key}: expected ${expectedType}, got ${typeof value}`);
        }
      }
    }

    // Apply validated updates
    Object.assign(ctx.state, updates, {
      lastUpdate: new Date().toISOString(),
      version: (ctx.state.version || 0) + 1
    });

    return { success: true, version: ctx.state.version };

  } catch (error) {
    // Restore on validation failure
    Object.assign(ctx.state, JSON.parse(backup));
    throw error;
  }
}

async function atomicStateUpdate(ctx: CallbackContext, operation: () => Promise<void>) {
  const stateSnapshot = JSON.stringify(ctx.state);
  const transactionId = `tx_${Date.now()}`;

  try {
    // Mark transaction start
    ctx.state._transaction = {
      id: transactionId,
      started: new Date().toISOString()
    };

    // Execute operation
    await operation();

    // Mark transaction complete
    ctx.state._transaction.completed = new Date().toISOString();
    delete ctx.state._transaction;

  } catch (error) {
    // Rollback on failure
    Object.assign(ctx.state, JSON.parse(stateSnapshot));

    // Log failed transaction
    ctx.state.lastFailedTransaction = {
      id: transactionId,
      error: error instanceof Error ? error.message : String(error),
      timestamp: new Date().toISOString()
    };

    throw error;
  }
}

async function resilientServiceOperation(ctx: ToolContext, query: string) {
  const results = {
    query,
    timestamp: new Date().toISOString(),
    sources: {
      memory: null as any,
      artifacts: null as any,
      state: ctx.state
    },
    fallbacksUsed: [] as string[]
  };

  // Try memory service first
  try {
    if (ctx.searchMemory) {
      results.sources.memory = await ctx.searchMemory(query);
    } else {
      results.fallbacksUsed.push("memory-unavailable");
    }
  } catch (error) {
    console.warn("Memory service failed:", error);
    results.fallbacksUsed.push("memory-failed");

    // Fallback to state-based search
    const stateKeys = Object.keys(ctx.state);
    const relevantKeys = stateKeys.filter(key =>
      key.toLowerCase().includes(query.toLowerCase())
    );
    results.sources.memory = {
      memories: relevantKeys.map(key => ({ content: `${key}: ${ctx.state[key]}` }))
    };
  }

  // Try artifact listing
  try {
    if (ctx.listArtifacts) {
      const artifacts = await ctx.listArtifacts();
      results.sources.artifacts = artifacts.filter(name =>
        name.toLowerCase().includes(query.toLowerCase())
      );
    } else {
      results.fallbacksUsed.push("artifacts-unavailable");
    }
  } catch (error) {
    console.warn("Artifact service failed:", error);
    results.fallbacksUsed.push("artifacts-failed");
    results.sources.artifacts = [];
  }

  return results;
}

async function coordinatedOperation(ctx: InvocationContext, operation: string) {
  const coordinationId = `coord_${Date.now()}`;
  const results = [];

  // Track coordination in session state
  ctx.session.state.activeCoordinations =
    ctx.session.state.activeCoordinations || {};
  ctx.session.state.activeCoordinations[coordinationId] = {
    operation,
    started: new Date().toISOString(),
    services: []
  };

  try {
    // Memory service operation
    if (ctx.memoryService) {
      const memoryResult = await ctx.memoryService.searchMemory({
        query: operation,
        appName: ctx.appName,
        userId: ctx.userId
      });

      results.push({ service: "memory", success: true, data: memoryResult });
      ctx.session.state.activeCoordinations[coordinationId].services.push("memory");
    }

    // Artifact service operation
    if (ctx.artifactService) {
      const artifacts = await ctx.artifactService.listArtifactKeys({
        appName: ctx.appName,
        userId: ctx.userId,
        sessionId: ctx.session.id
      });

      results.push({ service: "artifact", success: true, data: artifacts });
      ctx.session.state.activeCoordinations[coordinationId].services.push("artifact");
    }

    // Session service operation
    await ctx.sessionService.updateSession(ctx.session);
    results.push({ service: "session", success: true, data: "updated" });
    ctx.session.state.activeCoordinations[coordinationId].services.push("session");

    // Mark as completed
    ctx.session.state.activeCoordinations[coordinationId].completed = new Date().toISOString();

    return results;

  } catch (error) {
    // Handle coordination failure
    ctx.session.state.activeCoordinations[coordinationId].error = {
      message: error instanceof Error ? error.message : String(error),
      timestamp: new Date().toISOString()
    };
    throw error;
  }
}

async function recoverableOperation(ctx: CallbackContext | ToolContext) {
  const operationId = `op_${Date.now()}`;

  try {
    // Attempt primary operation
    const result = await performPrimaryOperation(ctx);

    // Save successful result
    await ctx.saveArtifact(`result_${operationId}.json`, {
      text: JSON.stringify(result, null, 2)
    });

    return result;

  } catch (primaryError) {
    console.warn("Primary operation failed:", primaryError);

    // Try recovery based on context capabilities
    if ('searchMemory' in ctx && ctx.searchMemory) {
      // ToolContext - try memory-based recovery
      try {
        const memoryResults = await ctx.searchMemory("similar operations");
        const recoveryData = memoryResults.memories?.[0];

        if (recoveryData) {
          ctx.state.recoveryUsed = {
            method: "memory",
            operationId,
            timestamp: new Date().toISOString()
          };

          return { recovered: true, data: recoveryData };
        }
      } catch (memoryError) {
        console.warn("Memory recovery failed:", memoryError);
      }
    }

    // Fallback to state-based recovery
    const fallbackData = ctx.state.lastSuccessfulOperation || null;
    if (fallbackData) {
      ctx.state.recoveryUsed = {
        method: "state-fallback",
        operationId,
        timestamp: new Date().toISOString()
      };

      return { recovered: true, data: fallbackData };
    }

    // No recovery possible - save error and rethrow
    ctx.state.lastOperationError = {
      operationId,
      error: primaryError instanceof Error ? primaryError.message : String(primaryError),
      timestamp: new Date().toISOString()
    };

    throw primaryError;
  }
}

async function performPrimaryOperation(ctx: CallbackContext | ToolContext) {
  // Placeholder for actual operation
  throw new Error("Simulated primary operation failure");
}

class ErrorHandlingAgent extends BaseAgent {
  protected async *runAsyncImpl(
    context: InvocationContext
  ): AsyncGenerator<Event, void, unknown> {
    const errorContext = {
      invocationId: context.invocationId,
      agentName: this.name,
      errors: [] as any[]
    };

    try {
      yield* this.processWithErrorHandling(context, errorContext);

    } catch (criticalError) {
      // Handle critical errors that affect the entire invocation
      yield this.createCriticalErrorEvent(context, criticalError, errorContext);

      // Attempt graceful degradation
      yield* this.attemptGracefulDegradation(context, errorContext);
    }
  }

  private async *processWithErrorHandling(
    context: InvocationContext,
    errorContext: any
  ): AsyncGenerator<Event, void, unknown> {
    // Try normal processing with error tracking
    try {
      // Update session state
      context.session.state.processingStatus = "active";

      // Process with services
      if (context.memoryService) {
        try {
          const memories = await context.memoryService.searchMemory({
            query: "user context",
            appName: context.appName,
            userId: context.userId
          });

          context.session.state.contextMemories = memories.memories?.length || 0;

        } catch (memoryError) {
          errorContext.errors.push({
            service: "memory",
            error: memoryError instanceof Error ? memoryError.message : String(memoryError),
            timestamp: new Date().toISOString()
          });

          // Continue without memory - non-critical error
        }
      }

      // Generate success event
      yield new Event({
        invocationId: context.invocationId,
        author: this.name,
        content: {
          parts: [{ text: "Processing completed with error handling" }]
        }
      });

    } catch (error) {
      errorContext.errors.push({
        stage: "processing",
        error: error instanceof Error ? error.message : String(error),
        timestamp: new Date().toISOString()
      });
      throw error;
    }
  }

  private createCriticalErrorEvent(
    context: InvocationContext,
    error: unknown,
    errorContext: any
  ): Event {
    return new Event({
      invocationId: context.invocationId,
      author: this.name,
      content: {
        parts: [{
          text: `Critical error occurred. Error details saved to session state.`
        }]
      },
      errorCode: "CRITICAL_ERROR",
      errorMessage: error instanceof Error ? error.message : String(error)
    });
  }
}

class CachingToolContext {
  private cache = new Map<string, { data: any; timestamp: number; ttl: number }>();

  constructor(private baseContext: ToolContext) {}

  async cachedSearchMemory(query: string, ttlMs = 60000): Promise<any> {
    const cacheKey = `memory_${query}`;
    const cached = this.cache.get(cacheKey);

    if (cached && Date.now() - cached.timestamp < cached.ttl) {
      // Update base context state to reflect cache hit
      this.baseContext.state.cacheHits = (this.baseContext.state.cacheHits || 0) + 1;
      return cached.data;
    }

    // Cache miss - perform actual search
    const result = await this.baseContext.searchMemory(query);

    // Cache the result
    this.cache.set(cacheKey, {
      data: result,
      timestamp: Date.now(),
      ttl: ttlMs
    });

    // Update context state
    this.baseContext.state.cacheMisses = (this.baseContext.state.cacheMisses || 0) + 1;

    return result;
  }

  async cachedListArtifacts(ttlMs = 30000): Promise<string[]> {
    const cacheKey = "artifacts_list";
    const cached = this.cache.get(cacheKey);

    if (cached && Date.now() - cached.timestamp < cached.ttl) {
      return cached.data;
    }

    const result = await this.baseContext.listArtifacts();

    this.cache.set(cacheKey, {
      data: result,
      timestamp: Date.now(),
      ttl: ttlMs
    });

    return result;
  }

  clearCache(): void {
    this.cache.clear();
    this.baseContext.state.cacheCleared = new Date().toISOString();
  }
}

async function batchContextOperations(ctx: ToolContext, operations: any[]) {
  const batchId = `batch_${Date.now()}`;
  const results = [];

  // Group operations by type
  const memoryQueries = operations.filter(op => op.type === "memory");
  const artifactOps = operations.filter(op => op.type === "artifact");
  const stateOps = operations.filter(op => op.type === "state");

  // Track batch in state
  ctx.state.activeBatches = ctx.state.activeBatches || {};
  ctx.state.activeBatches[batchId] = {
    started: new Date().toISOString(),
    totalOps: operations.length,
    completed: 0
  };

  try {
    // Batch memory queries
    if (memoryQueries.length > 0) {
      const memoryPromises = memoryQueries.map(async (op) => {
        try {
          const result = await ctx.searchMemory(op.query);
          ctx.state.activeBatches[batchId].completed++;
          return { operation: op, result, success: true };
        } catch (error) {
          return { operation: op, error: String(error), success: false };
        }
      });

      const memoryResults = await Promise.all(memoryPromises);
      results.push(...memoryResults);
    }

    // Batch artifact operations
    if (artifactOps.length > 0) {
      // Batch artifact loading
      const artifactPromises = artifactOps.map(async (op) => {
        try {
          let result;
          if (op.action === "list") {
            result = await ctx.listArtifacts();
          } else if (op.action === "load") {
            result = await ctx.loadArtifact(op.filename);
          }

          ctx.state.activeBatches[batchId].completed++;
          return { operation: op, result, success: true };
        } catch (error) {
          return { operation: op, error: String(error), success: false };
        }
      });

      const artifactResults = await Promise.all(artifactPromises);
      results.push(...artifactResults);
    }

    // Batch state operations (synchronous)
    stateOps.forEach(op => {
      try {
        if (op.action === "set") {
          ctx.state[op.key] = op.value;
        } else if (op.action === "get") {
          results.push({
            operation: op,
            result: ctx.state[op.key],
            success: true
          });
        }
        ctx.state.activeBatches[batchId].completed++;
      } catch (error) {
        results.push({
          operation: op,
          error: String(error),
          success: false
        });
      }
    });

    // Mark batch complete
    ctx.state.activeBatches[batchId].completedAt = new Date().toISOString();

    return {
      batchId,
      total: operations.length,
      successful: results.filter(r => r.success).length,
      failed: results.filter(r => !r.success).length,
      results
    };

  } catch (error) {
    ctx.state.activeBatches[batchId].error = {
      message: error instanceof Error ? error.message : String(error),
      timestamp: new Date().toISOString()
    };
    throw error;
  }
}

// ❌ DON'T: Use powerful context for simple operations
function badReadOnlyOperation(ctx: ToolContext) {
  // Only reading state but using ToolContext unnecessarily
  return ctx.state.userName;
}

// ✅ DO: Use appropriate context level
function goodReadOnlyOperation(ctx: ReadonlyContext) {
  return ctx.state.userName;
}

// ❌ DON'T: Hold references to context beyond operation scope
class BadService {
  private storedContext: ToolContext; // Don't do this!

  setContext(ctx: ToolContext) {
    this.storedContext = ctx; // Context may become stale
  }
}

// ✅ DO: Pass context to each operation
class GoodService {
  async performOperation(ctx: ToolContext, params: any) {
    // Use context immediately, don't store it
    return await ctx.searchMemory(params.query);
  }
}

// ❌ DON'T: Ignore service availability
async function badServiceUsage(ctx: ToolContext) {
  // Will throw if memory service not configured
  return await ctx.searchMemory("query");
}

// ✅ DO: Check service availability
async function goodServiceUsage(ctx: ToolContext) {
  try {
    return await ctx.searchMemory("query");
  } catch (error) {
    if (error.message.includes("Memory service is not available")) {
      // Handle gracefully
      return { memories: [] };
    }
    throw error;
  }
}

// Create mock contexts for testing
function createMockReadonlyContext(state: any = {}): ReadonlyContext {
  return {
    userContent: { parts: [{ text: "test input" }] },
    invocationId: "test-invocation-123",
    agentName: "test-agent",
    state: Object.freeze(state)
  } as ReadonlyContext;
}

function createMockToolContext(state: any = {}): ToolContext {
  return {
    ...createMockReadonlyContext(state),
    state: state, // Mutable for ToolContext
    functionCallId: "test-function-call-456",
    actions: {
      escalate: false,
      transferToAgent: null,
      skipSummarization: false,
      stateDelta: {},
      artifactDelta: {}
    },
    loadArtifact: jest.fn(),
    saveArtifact: jest.fn(),
    listArtifacts: jest.fn().mockResolvedValue([]),
    searchMemory: jest.fn().mockResolvedValue({ memories: [] })
  } as any;
}

// Test context usage
describe("Context Usage", () => {
  test("should handle readonly context safely", () => {
    const mockState = { user: { name: "Test User" } };
    const ctx = createMockReadonlyContext(mockState);

    const result = analyzeState(ctx);

    expect(result.userCount).toBeGreaterThan(0);
    expect(ctx.state).toEqual(mockState); // State unchanged
  });

  test("should use tool context features", async () => {
    const ctx = createMockToolContext({ preferences: { theme: "dark" } });

    const result = await smartTool({ query: "test" }, ctx);

    expect(ctx.searchMemory).toHaveBeenCalledWith("test");
    expect(ctx.state.lastSearch).toBeDefined();
  });
});