index.ts

import express from 'express';
import { v4 as uuidv4 } from 'uuid'; // For generating unique IDs

import {
  A2A_PROTOCOL_VERSION,
  AgentCard,
  Task,
  TaskState,
  TaskStatusUpdateEvent,
  Message,
  AGENT_CARD_PATH,
  TaskArtifactUpdateEvent,
  Artifact,
  Part,
  Role,
} from '../../../index.js';
import {
  InMemoryTaskStore,
  TaskStore,
  AgentExecutor,
  RequestContext,
  ExecutionEventBus,
  DefaultRequestHandler,
  AgentEvent,
} from '../../../server/index.js';
import { agentCardHandler, jsonRpcHandler, UserBuilder } from '../../../server/express/index.js';
import { MessageData } from 'genkit';
import { ai } from './genkit.js';
import { searchMovies, searchPeople } from './tools.js';

if (!process.env.GEMINI_API_KEY || !process.env.TMDB_API_KEY) {
  console.error('GEMINI_API_KEY and TMDB_API_KEY environment variables are required');
  process.exit(1);
}

// Simple store for contexts
const contexts: Map<string, Message[]> = new Map();

// Load the Genkit prompt
const movieAgentPrompt = ai.prompt('movie_agent');

/**
 * MovieAgentExecutor implements the agent's core logic.
 */
class MovieAgentExecutor implements AgentExecutor {
  private cancelledTasks = new Set<string>();

  public cancelTask = async (taskId: string, _eventBus: ExecutionEventBus): Promise<void> => {
    this.cancelledTasks.add(taskId);
    // The execute loop is responsible for publishing the final state
  };

  async execute(requestContext: RequestContext, eventBus: ExecutionEventBus): Promise<void> {
    const userMessage = requestContext.userMessage;
    const existingTask = requestContext.task;

    // Determine IDs for the task and context
    const taskId = requestContext.taskId;
    const contextId = requestContext.contextId;

    console.log(
      `[MovieAgentExecutor] Processing message ${userMessage.messageId} for task ${taskId} (context: ${contextId})`
    );

    // 1. Every streaming turn must begin with a Task or Message event.
    const taskSnapshot: Task = existingTask ?? {
      id: taskId,
      contextId: contextId,
      status: {
        state: TaskState.TASK_STATE_SUBMITTED,
        timestamp: new Date().toISOString(),
        message: undefined,
      },
      artifacts: [],
      history: [userMessage],
      metadata: userMessage.metadata,
    };
    eventBus.publish(AgentEvent.task(taskSnapshot));

    // 2. Publish "working" status update
    const workingStatusUpdate: TaskStatusUpdateEvent = {
      taskId: taskId,
      contextId: contextId,
      status: {
        state: TaskState.TASK_STATE_WORKING,
        message: {
          role: Role.ROLE_AGENT,
          messageId: uuidv4(),
          parts: [
            {
              content: { $case: 'text', value: 'Processing your question, hang tight!' },
              metadata: undefined,
              filename: '',
              mediaType: 'text/plain',
            },
          ],
          taskId: taskId,
          contextId: contextId,
          extensions: [],
          metadata: {},
          referenceTaskIds: [],
        },
        timestamp: new Date().toISOString(),
      },
      metadata: {},
    };
    eventBus.publish(AgentEvent.statusUpdate(workingStatusUpdate));

    // 3. Prepare messages for Genkit prompt
    const historyForGenkit = contexts.get(contextId) || [];
    if (!historyForGenkit.find((m) => m.messageId === userMessage.messageId)) {
      historyForGenkit.push(userMessage);
    }
    contexts.set(contextId, historyForGenkit);

    const messages: MessageData[] = historyForGenkit
      .map((m) => {
        const textContent = m.parts
          .map((p) => (p.content?.$case === 'text' ? p.content.value : ''))
          .filter((t) => !!t)
          .join('\n');

        return {
          role: (m.role === Role.ROLE_AGENT ? 'model' : 'user') as 'user' | 'model',
          content: textContent ? [{ text: textContent }] : [],
        };
      })
      .filter((m) => m.content.length > 0);

    if (messages.length === 0) {
      console.warn(
        `[MovieAgentExecutor] No valid text messages found in history for task ${taskId}.`
      );
      const failureUpdate: TaskStatusUpdateEvent = {
        taskId: taskId,
        contextId: contextId,
        status: {
          state: TaskState.TASK_STATE_FAILED,
          message: {
            role: Role.ROLE_AGENT,
            messageId: uuidv4(),
            parts: [
              {
                content: { $case: 'text', value: 'No message found to process.' },
                metadata: undefined,
                filename: '',
                mediaType: 'text/plain',
              },
            ],
            taskId: taskId,
            contextId: contextId,
            extensions: [],
            metadata: {},
            referenceTaskIds: [],
          },
          timestamp: new Date().toISOString(),
        },
        metadata: {},
      };
      eventBus.publish(AgentEvent.statusUpdate(failureUpdate));
      return;
    }

    const goal =
      (existingTask?.metadata?.goal as string | undefined) ||
      (userMessage.metadata?.goal as string | undefined);

    try {
      // 4. Run the Genkit prompt
      const response = await movieAgentPrompt(
        { goal: goal, now: new Date().toISOString() },
        {
          messages,
          tools: [searchMovies, searchPeople],
        }
      );

      // Check if the request has been cancelled
      if (this.cancelledTasks.has(taskId)) {
        console.log(`[MovieAgentExecutor] Request cancelled for task: ${taskId}`);

        const cancelledUpdate: TaskStatusUpdateEvent = {
          taskId: taskId,
          contextId: contextId,
          status: {
            state: TaskState.TASK_STATE_CANCELED,
            timestamp: new Date().toISOString(),
            message: undefined,
          },
          metadata: {},
        };
        eventBus.publish(AgentEvent.statusUpdate(cancelledUpdate));
        return;
      }

      const responseText = response.text;
      console.info(`[MovieAgentExecutor] Prompt response: ${responseText}`);
      const lines = responseText.trim().split('\n');
      const finalStateLine = lines.at(-1)?.trim().toUpperCase();
      const agentReplyText = lines
        .slice(0, lines.length - 1)
        .join('\n')
        .trim();

      let finalA2AState: TaskState = TaskState.TASK_STATE_UNSPECIFIED;

      if (finalStateLine === 'COMPLETED') {
        finalA2AState = TaskState.TASK_STATE_COMPLETED;
      } else if (finalStateLine === 'AWAITING_USER_INPUT') {
        finalA2AState = TaskState.TASK_STATE_INPUT_REQUIRED;
      } else {
        console.warn(
          `[MovieAgentExecutor] Unexpected final state line from prompt: "${finalStateLine}". Defaulting to 'completed'.`
        );
        finalA2AState = TaskState.TASK_STATE_COMPLETED;
      }

      // 5. Publish artifact with the result
      const parts: Part[] = [
        {
          content: { $case: 'text', value: agentReplyText || 'Completed.' },
          metadata: undefined,
          filename: '',
          mediaType: 'text/plain',
        },
      ];
      const artifactId = uuidv4();
      const resultArtifact: Artifact = {
        artifactId: artifactId,
        name: 'Result',
        description: 'The result of the movie agent.',
        parts: parts,
        metadata: undefined,
        extensions: [],
      };

      const artifactUpdate: TaskArtifactUpdateEvent = {
        taskId: taskId,
        contextId: contextId,
        artifact: resultArtifact,
        lastChunk: true,
        append: false,
        metadata: {},
      };
      eventBus.publish(AgentEvent.artifactUpdate(artifactUpdate));

      // 6. Update local history context (internal only)
      const agentMessage: Message = {
        role: Role.ROLE_AGENT,
        messageId: uuidv4(),
        parts: parts,
        taskId: taskId,
        contextId: contextId,
        extensions: [],
        metadata: {},
        referenceTaskIds: [],
      };
      historyForGenkit.push(agentMessage);
      contexts.set(contextId, historyForGenkit);

      // 7. Publish final task status update
      const finalUpdate: TaskStatusUpdateEvent = {
        taskId: taskId,
        contextId: contextId,
        status: {
          state: finalA2AState,
          timestamp: new Date().toISOString(),
          message: undefined,
        },
        metadata: {},
      };
      eventBus.publish(AgentEvent.statusUpdate(finalUpdate));

      console.log(`[MovieAgentExecutor] Task ${taskId} finished with state: ${finalA2AState}`);
    } catch (error: any) {
      console.error(`[MovieAgentExecutor] Error processing task ${taskId}:`, error);
      const errorUpdate: TaskStatusUpdateEvent = {
        taskId: taskId,
        contextId: contextId,
        status: {
          state: TaskState.TASK_STATE_FAILED,
          message: {
            role: Role.ROLE_AGENT,
            messageId: uuidv4(),
            parts: [
              {
                content: { $case: 'text', value: `Agent error: ${error.message}` },
                metadata: undefined,
                filename: '',
                mediaType: 'text/plain',
              },
            ],
            taskId: taskId,
            contextId: contextId,
            extensions: [],
            metadata: undefined,
            referenceTaskIds: [],
          },
          timestamp: new Date().toISOString(),
        },
        metadata: undefined,
      };
      eventBus.publish(AgentEvent.statusUpdate(errorUpdate));
    }
  }
}

// --- Server Setup ---

const movieAgentCard: AgentCard = {
  name: 'Movie Agent',
  description: 'An agent that can answer questions about movies and actors using TMDB.',
  supportedInterfaces: [
    {
      url: 'http://localhost:41241/',
      protocolBinding: 'JSONRPC',
      tenant: '',
      protocolVersion: A2A_PROTOCOL_VERSION,
    },
  ],
  provider: {
    organization: 'A2A Samples',
    url: 'https://example.com/a2a-samples',
  },
  version: '0.0.2',
  capabilities: {
    streaming: true,
    pushNotifications: false,
    extensions: [],
    extendedAgentCard: false,
  },
  securitySchemes: {}, // Or define actual security schemes if any
  securityRequirements: [],
  defaultInputModes: ['text'],
  defaultOutputModes: ['text', 'task-status'], // task-status is a common output mode
  skills: [
    {
      id: 'general_movie_chat',
      name: 'General Movie Chat',
      description: 'Answer general questions or chat about movies, actors, directors.',
      tags: ['movies', 'actors', 'directors'],
      examples: [
        'Tell me about the plot of Inception.',
        'Recommend a good sci-fi movie.',
        'Who directed The Matrix?',
        'What other movies has Scarlett Johansson been in?',
        'Find action movies starring Keanu Reeves',
        'Which came out first, Jurassic Park or Terminator 2?',
      ],
      inputModes: ['text'], // Explicitly defining for skill
      outputModes: ['text', 'task-status'], // Explicitly defining for skill
      securityRequirements: [],
    },
  ],
  documentationUrl: '',
  signatures: [],
};

async function main() {
  // 1. Create TaskStore
  const taskStore: TaskStore = new InMemoryTaskStore();

  // 2. Create AgentExecutor
  const agentExecutor: AgentExecutor = new MovieAgentExecutor();

  // 3. Create DefaultRequestHandler
  const requestHandler = new DefaultRequestHandler(movieAgentCard, taskStore, agentExecutor);

  // 4. Create and setup Express.js app
  const app = express();

  app.use(`/${AGENT_CARD_PATH}`, agentCardHandler({ agentCardProvider: requestHandler }));
  app.use(jsonRpcHandler({ requestHandler, userBuilder: UserBuilder.noAuthentication }));

  // 5. Start the server
  const PORT = process.env.PORT || 41241;
  app.listen(PORT, (err) => {
    if (err) {
      throw err;
    }
    console.log(`[MovieAgent] Server using new framework started on http://localhost:${PORT}`);
    console.log(`[MovieAgent] Agent Card: http://localhost:${PORT}/.well-known/agent-card.json`);
    console.log('[MovieAgent] Press Ctrl+C to stop the server');
  });
}

main().catch(console.error);