// src/agent.ts
import { AriaAgent } from "@projectaria/cf-agents";
import { createAgent } from "@projectaria/aria-agents";
import { openai } from "@ai-sdk/openai";

export class MyAssistant extends AriaAgent {
  // Create an ARIA agent with your configuration
  private ariaAgent = createAgent({
    id: "assistant",
    name: "My Assistant",
    instructions: "You are a helpful assistant that remembers user preferences.",
    model: openai("gpt-4o"),
  });

  // Handle incoming chat messages
  async onChatMessage(onFinish: () => void) {
    // Get the current session with message history
    const session = this.getSession();
    const lastMessage = session.messages[session.messages.length - 1];

    if (!lastMessage) {
      return;
    }

    // Run the ARIA agent with automatic memory persistence
    const result = await this.runAriaAgent(this.ariaAgent, lastMessage.content);

    // Call onFinish to signal completion
    onFinish();

    // Return response (for HTTP) or it's already sent via WebSocket
    return new Response(result.text);
  }
}

// src/index.ts
import { routeAgentRequest } from "agents";
import { MyAssistant } from "./agent";

export { MyAssistant };

export default {
  async fetch(request: Request, env: Env): Promise<Response> {
    // Route requests to your agent
    return (
      (await routeAgentRequest(request, env)) ||
      new Response("Not found", { status: 404 })
    );
  },
};

class MyAgent extends AriaAgent {
  async onChatMessage(onFinish) {
    // Store data for current session
    this.setWorkingMemory("cart", { items: ["item1", "item2"] });
    
    // Retrieve working memory
    const memory = this.getWorkingMemory();
    console.log(memory.cart); // { items: ["item1", "item2"] }
  }
}

class MyAgent extends AriaAgent {
  async onChatMessage(onFinish) {
    // Store a preference that persists across sessions
    this.storeMemory({
      key: "user:theme",
      value: { theme: "dark", fontSize: "large" },
      type: "preference",
    });

    // Search memories
    const results = this.searchMemories({
      query: "theme",
      types: ["preference"],
      limit: 5,
    });
  }
}

import { createFlow } from "@projectaria/aria-agents";

class MyAgent extends AriaAgent {
  private orderFlow = createFlow({
    id: "order-flow",
    nodes: [
      { id: "collect-items", ... },
      { id: "confirm-order", ... },
      { id: "payment", ... },
    ],
  });

  constructor(ctx, env) {
    super(ctx, env);
    this.registerFlow(this.orderFlow);
  }
}

class AriaAgent extends AIChatAgent<AriaEnv, AriaSessionState> {
  // Session Management
  getSession(): AriaSession
  
  // Working Memory
  getWorkingMemory(): WorkingMemory
  setWorkingMemory(key: string, value: unknown): void
  clearWorkingMemory(): void
  
  // Cross-Session Memory
  storeMemory(memory: Omit<AriaMemory, "id" | ...>): string
  searchMemories(query: MemorySearchQuery): MemorySearchResult[]
  
  // ARIA Integration
  runAriaAgent(agent, input, options?): Promise<{ text: string; session: AgentSession }>
  streamAriaAgent(agent, input, options?): AsyncIterable<string>
  
  // Flow Orchestration
  registerFlow(flow: FlowDefinition): void
  transitionFlow(flowId: string, toNodeId: string): void
}

const ws = new WebSocket("wss://your-worker.workers.dev/agents/my-assistant/session-123");

ws.onopen = () => {
  ws.send(JSON.stringify({
    type: "cf.agent.chat.request",
    init: {
      method: "POST",
      body: JSON.stringify({
        messages: [{ role: "user", content: "Hello!" }]
      })
    }
  }));
};

ws.onmessage = (event) => {
  const data = JSON.parse(event.data);
  console.log("Response:", data);
};

const response = await fetch("https://your-worker.workers.dev/agents/my-assistant/session-123", {
  method: "POST",
  headers: { "Content-Type": "application/json" },
  body: JSON.stringify({
    messages: [{ role: "user", content: "Hello!" }]
  })
});

// src/types.ts
interface Env {
  MyAssistant: DurableObjectNamespace;
  OPENAI_API_KEY: string;
}

// src/shopping-assistant.ts
import { AriaAgent } from "@projectaria/cf-agents";
import { createAgent } from "@projectaria/aria-agents";
import { openai } from "@ai-sdk/openai";
import { tool } from "ai";
import { z } from "zod";

export class ShoppingAssistant extends AriaAgent {
  private agent = createAgent({
    id: "shopping-assistant",
    name: "Shopping Assistant",
    instructions: `You are a helpful shopping assistant.
You remember user preferences and can help them find products.
Use the cart tools to manage their shopping cart.`,
    model: openai("gpt-4o"),
    tools: {
      addToCart: tool({
        description: "Add an item to the cart",
        parameters: z.object({
          item: z.string(),
          quantity: z.number().default(1),
        }),
        execute: async ({ item, quantity }) => {
          const cart = this.getWorkingMemory().cart || [];
          cart.push({ item, quantity });
          this.setWorkingMemory("cart", cart);
          return { success: true, cart };
        },
      }),
      viewCart: tool({
        description: "View the current cart",
        parameters: z.object({}),
        execute: async () => {
          return this.getWorkingMemory().cart || [];
        },
      }),
    },
  });

  async onChatMessage(onFinish: () => void) {
    const session = this.getSession();
    const lastMessage = session.messages[session.messages.length - 1];

    if (!lastMessage) return;

    // Check for returning user preferences
    const prefs = this.searchMemories({
      types: ["preference"],
      limit: 5,
    });

    // Run agent with context
    const result = await this.runAriaAgent(this.agent, lastMessage.content);

    onFinish();
    return new Response(result.text);
  }
}