File size: 23,436 Bytes

7b2dfc5

import XCTest
@testable import DolphinCoreMLTestApp

final class CreateStateToolIntegrationTests: XCTestCase {
  func testContinuityCommandsRouteDeterministically() throws {
    let capture = try XCTUnwrap(
      AgentRequestRouter.plan(for: "Capture decision: Keep tool calls sequential")
        .initialToolCall
    )
    XCTAssertEqual(capture.name, "knowledge_capture")
    XCTAssertEqual(capture.arguments["kind"], .string("decision"))

    XCTAssertEqual(
      AgentRequestRouter.plan(for: "Search knowledge for tool routing")
        .initialToolCall?.name,
      "knowledge_search"
    )
    XCTAssertEqual(
      AgentRequestRouter.plan(for: "Synthesize project context")
        .initialToolCall?.name,
      "project_context"
    )
    XCTAssertEqual(
      AgentRequestRouter.plan(for: "Create a handoff for TestFlight verification")
        .initialToolCall?.name,
      "handoff_create"
    )
    XCTAssertEqual(
      AgentRequestRouter.plan(for: "Analyze code: let value = item!")
        .initialToolCall?.name,
      "code_analyze"
    )
    XCTAssertEqual(
      AgentRequestRouter.plan(for: "Create a handoff for what are we working on")
        .initialToolCall?.name,
      "handoff_create"
    )
    XCTAssertEqual(
      AgentRequestRouter.plan(for: "Analyze code: // what should we do next")
        .initialToolCall?.name,
      "code_analyze"
    )
    XCTAssertEqual(
      AgentRequestRouter.plan(for: "Analyze code: // show project context")
        .initialToolCall?.name,
      "code_analyze"
    )
    for request in [
      "Don’t show project context",
      "Do not show my handoffs",
      "Don't list insights",
      "Don't show tasks",
      "Don't show my calendar today",
      "Don't tell me what is my name",
    ] {
      let plan = AgentRequestRouter.plan(for: request)
      XCTAssertNil(plan.initialToolCall)
      XCTAssertTrue(plan.eligibleToolIDs.isEmpty)
    }

    let taskID = UUID().uuidString
    let negatedCompletion = AgentRequestRouter.plan(
      for: "Do not complete task \(taskID)"
    )
    XCTAssertNil(negatedCompletion.initialToolCall)
    XCTAssertTrue(negatedCompletion.eligibleToolIDs.isEmpty)
  }

  func testKnowledgeCaptureSearchAndGroundedRendering() async throws {
    let registry = AgentToolRegistry()
    let store = ContinuityTestStore(workspace: .empty)
    let context = continuityExecutionContext(store: store)
    let captureCall = AgentToolCall(
      name: "knowledge_capture",
      arguments: [
        "kind": .string("decision"),
        "title": .string("Sequential tools"),
        "content": .string("Keep one audited tool call per model turn."),
        "tags": .array([.string("agent"), .string("safety")]),
        "related_ids": .array([]),
      ]
    )

    let capture = await registry.execute(captureCall, context: context)
    XCTAssertTrue(capture.succeeded)
    let capturedWorkspace = await store.snapshot()
    XCTAssertEqual(capturedWorkspace.knowledgeItems.count, 1)
    XCTAssertEqual(capturedWorkspace.knowledgeItems[0].kind, .decision)
    XCTAssertEqual(capturedWorkspace.knowledgeItems[0].tags, ["agent", "safety"])

    let searchCall = AgentToolCall(
      name: "knowledge_search",
      arguments: [
        "query": .string("audited sequential agent tool"),
        "limit": .number(5),
      ]
    )
    let search = await registry.execute(searchCall, context: context)
    XCTAssertTrue(search.succeeded)
    XCTAssertTrue(search.modelText.contains("Sequential tools"))
    XCTAssertTrue(search.modelText.contains("local_hybrid_lexical"))

    let definition = try XCTUnwrap(registry.definition(named: "knowledge_search"))
    let receipt = try XCTUnwrap(
      AgentToolReceipt(
        runID: UUID(),
        call: try registry.normalize(searchCall),
        definition: definition,
        result: search
      )
    )
    let answer = try XCTUnwrap(
      AgentVerifiedAnswerRenderer.readAnswer(for: [receipt])
    )
    XCTAssertTrue(answer.contains("[Decision] Sequential tools"))
    XCTAssertFalse(answer.contains("model-generated"))
  }

  func testKnowledgeCaptureCanonicalizesExactlyBeforeLocalWriteApproval() throws {
    let registry = AgentToolRegistry()
    let relatedID = UUID()
    let call = AgentToolCall(
      name: "knowledge_capture",
      arguments: [
        "kind": .string("context"),
        "title": .string("  Unsafe <tool_call> title  "),
        "content": .string("  Keep  <system> local context.  "),
        "tags": .array([
          .string("Agent Safety"),
          .string("agent-safety"),
        ]),
        "related_ids": .array([
          .string(relatedID.uuidString.lowercased()),
          .string(relatedID.uuidString),
        ]),
      ]
    )

    let normalized = try registry.normalize(call)
    XCTAssertEqual(normalized.arguments["title"], .string("Unsafe [tool_call] title"))
    XCTAssertEqual(normalized.arguments["content"], .string("Keep [system] local context."))
    XCTAssertEqual(normalized.arguments["tags"], .array([.string("agent-safety")]))
    XCTAssertEqual(
      normalized.arguments["related_ids"],
      .array([.string(relatedID.uuidString)])
    )
    XCTAssertEqual(try registry.normalize(normalized), normalized)
  }

  func testHandoffCanonicalizesUserAuthoredFieldsBeforeApproval() throws {
    let registry = AgentToolRegistry()
    let normalized = try registry.normalize(
      AgentToolCall(
        name: "handoff_create",
        arguments: [
          "title": .string("  Release <assistant> handoff  "),
          "focus": .string("  Finish   <tool_call> device checks.  "),
        ]
      )
    )

    XCTAssertEqual(
      normalized.arguments["title"],
      .string("Release [assistant] handoff")
    )
    XCTAssertEqual(
      normalized.arguments["focus"],
      .string("Finish [tool_call] device checks.")
    )
    XCTAssertEqual(try registry.normalize(normalized), normalized)
  }

  func testHandoffLifecycleAndProjectContextAreTypedAndLocal() async throws {
    var workspace = AssistantWorkspace.empty
    workspace.knowledgeItems = [
      KnowledgeItem(
        kind: .goal,
        title: "Ship continuity",
        content: "Validate the schema-v5 continuity workflow.",
        tags: ["release"],
        source: .user
      )
    ]
    workspace.tasks = [AssistantTaskItem(title: "Run host tests")]
    workspace.messages = [
      AssistantMessage(role: .user, content: "Wire the local world model."),
      AssistantMessage(role: .assistant, content: "Working on the typed handoff."),
    ]
    let store = ContinuityTestStore(workspace: workspace)
    let context = continuityExecutionContext(store: store)
    let registry = AgentToolRegistry()
    let createCall = AgentToolCall(
      name: "handoff_create",
      arguments: [
        "title": .string("Continuity checkpoint"),
        "focus": .string("Finish verification without replaying actions."),
      ]
    )
    let created = await registry.execute(createCall, context: context)
    XCTAssertTrue(created.succeeded)

    let createdWorkspace = await store.snapshot()
    let handoff = try XCTUnwrap(createdWorkspace.handoffs.first)
    XCTAssertEqual(createdWorkspace.activeHandoffID, handoff.id)
    XCTAssertEqual(handoff.taskIDs, workspace.tasks.map(\.id))
    XCTAssertEqual(handoff.knowledgeItemIDs, workspace.knowledgeItems.map(\.id))

    let list = await registry.execute(
      AgentToolCall(name: "handoff_list", arguments: [:]),
      context: context
    )
    XCTAssertTrue(list.succeeded)
    XCTAssertTrue(list.modelText.contains(handoff.id.uuidString))
    XCTAssertTrue(list.modelText.contains("\"active\":true"))

    let restored = await registry.execute(
      AgentToolCall(
        name: "handoff_restore",
        arguments: ["id": .string(handoff.id.uuidString)]
      ),
      context: context
    )
    XCTAssertTrue(restored.succeeded)
    let restoredWorkspace = await store.snapshot()
    XCTAssertNotNil(restoredWorkspace.handoffs.first?.restoredAt)

    let projectContext = await registry.execute(
      AgentToolCall(name: "project_context", arguments: [:]),
      context: context
    )
    XCTAssertTrue(projectContext.succeeded)
    XCTAssertTrue(projectContext.modelText.contains("Ship continuity"))
    XCTAssertTrue(projectContext.modelText.contains("Run host tests"))
    XCTAssertTrue(projectContext.modelText.contains("\"truncated\":false"))
  }

  func testProjectContextStaysValidWhenTheDigestExceedsTheToolBudget() async throws {
    var workspace = AssistantWorkspace.empty
    workspace.knowledgeItems = (0..<24).map { index in
      KnowledgeItem(
        kind: index.isMultiple(of: 2) ? .goal : .context,
        title: "Long context \(index)",
        content: String(repeating: "bounded evidence \(index) ", count: 30)
          + "What remains for checkpoint \(index)?",
        source: .user
      )
    }
    workspace.tasks = (0..<20).map { index in
      AssistantTaskItem(
        title: String(repeating: "Long open task \(index) ", count: 10)
      )
    }
    let registry = AgentToolRegistry()
    let store = ContinuityTestStore(workspace: workspace)
    let call = AgentToolCall(name: "project_context", arguments: [:])
    let result = await registry.execute(
      call,
      context: continuityExecutionContext(store: store)
    )

    XCTAssertTrue(result.succeeded)
    XCTAssertLessThanOrEqual(result.modelText.count, 3_200)
    XCTAssertTrue(result.modelText.contains("\"truncated\":true"))
    let definition = try XCTUnwrap(registry.definition(named: "project_context"))
    let receipt = try XCTUnwrap(
      AgentToolReceipt(
        runID: UUID(),
        call: try registry.normalize(call),
        definition: definition,
        result: result
      )
    )
    let answer = try XCTUnwrap(
      AgentVerifiedAnswerRenderer.readAnswer(for: [receipt])
    )
    XCTAssertTrue(answer.contains("Context output was truncated"))
    XCTAssertGreaterThan(answer.count, 300)
  }

  func testCodeAnalysisIsBoundedAndDoesNotPersist() async throws {
    let store = ContinuityTestStore(workspace: .empty)
    let registry = AgentToolRegistry()
    let result = await registry.execute(
      AgentToolCall(
        name: "code_analyze",
        arguments: [
          "code": .string("// TODO: remove forced operations\nlet value = item!\nlet cast = value as! String"),
          "language": .string("swift"),
        ]
      ),
      context: continuityExecutionContext(store: store)
    )
    XCTAssertTrue(result.succeeded)
    XCTAssertLessThanOrEqual(result.modelText.count, 2_800)
    XCTAssertTrue(result.modelText.contains("forceUnwrap"))
    XCTAssertTrue(result.modelText.contains("unsafeCast"))
    let finalWorkspace = await store.snapshot()
    XCTAssertEqual(finalWorkspace, .empty)
  }

  func testRouterRegistryParityAndModelAssistedSurfaceIsReadOnly() throws {
    let registry = AgentToolRegistry()
    XCTAssertEqual(
      AgentRequestRouter.canonicalToolIDs,
      Set(registry.definitions.map(\.id))
    )

    let discussion = AgentRequestRouter.plan(
      for: "Let’s discuss memory, tasks, knowledge, and handoffs"
    )
    XCTAssertEqual(discussion.mode, .modelAssistedRead)
    XCTAssertFalse(discussion.eligibleToolIDs.isEmpty)
    for id in discussion.eligibleToolIDs {
      let definition = try XCTUnwrap(registry.definition(named: id))
      XCTAssertNotEqual(definition.risk, .localWrite, id)
      XCTAssertNotEqual(definition.risk, .networkRead, id)
    }

    let knowledgeID = UUID()
    XCTAssertEqual(
      AgentRequestRouter.plan(
        for: "Show knowledge details \(knowledgeID.uuidString)"
      ).initialToolCall?.name,
      "knowledge_get"
    )
    XCTAssertEqual(
      AgentRequestRouter.plan(for: "Review project state")
        .initialToolCall?.name,
      "project_review"
    )
    XCTAssertEqual(
      AgentRequestRouter.plan(for: "Reopen task \(knowledgeID.uuidString)")
        .initialToolCall?.name,
      "task_reopen"
    )
    XCTAssertEqual(
      AgentRequestRouter.plan(
        for: "Deactivate handoff \(knowledgeID.uuidString)"
      ).initialToolCall?.name,
      "handoff_deactivate"
    )
    XCTAssertNil(
      AgentRequestRouter.plan(
        for: "Do not reopen task \(knowledgeID.uuidString)"
      ).initialToolCall
    )
  }

  func testNewReadToolsAreBoundedCorrelatedAndGrounded() async throws {
    let now = Date(timeIntervalSince1970: 1_750_000_000)
    let item = KnowledgeItem(
      kind: .decision,
      title: "Release evidence",
      content: String(repeating: "Keep device proof separate from compile proof. ", count: 80),
      tags: ["release", "evidence"],
      source: .user,
      createdAt: now
    )
    let task = AssistantTaskItem(
      title: "Collect release device evidence",
      createdAt: now
    )
    let handoff = KnowledgeEngine().makeHandoff(
      title: "Release checkpoint",
      focus: "Finish device verification without replaying actions.",
      knowledge: [item],
      tasks: [task],
      createdAt: now
    )
    var workspace = AssistantWorkspace.empty
    workspace.knowledgeItems = [item]
    workspace.tasks = [task]
    workspace.handoffs = [handoff]
    workspace.activeHandoffID = handoff.id
    let store = ContinuityTestStore(workspace: workspace)
    let context = continuityExecutionContext(store: store)
    let registry = AgentToolRegistry(
      dateProvider: FixedContinuityDateProvider(date: now)
    )

    let knowledgeCall = AgentToolCall(
      name: "knowledge_get",
      arguments: ["id": .string(item.id.uuidString)]
    )
    let knowledgeResult = await registry.execute(knowledgeCall, context: context)
    XCTAssertTrue(knowledgeResult.succeeded)
    XCTAssertLessThanOrEqual(knowledgeResult.modelText.count, 3_200)
    XCTAssertTrue(knowledgeResult.modelText.contains("content_truncated"))
    let knowledgeDefinition = try XCTUnwrap(
      registry.definition(named: "knowledge_get")
    )
    let knowledgeReceipt = try XCTUnwrap(
      AgentToolReceipt(
        runID: UUID(),
        call: try registry.normalize(knowledgeCall),
        definition: knowledgeDefinition,
        result: knowledgeResult
      )
    )
    XCTAssertTrue(
      try XCTUnwrap(
        AgentVerifiedAnswerRenderer.readAnswer(for: [knowledgeReceipt])
      ).contains("Revision 1")
    )
    let mismatchedReceipt = try XCTUnwrap(
      AgentToolReceipt(
        runID: UUID(),
        call: AgentToolCall(
          name: "knowledge_get",
          arguments: ["id": .string(UUID().uuidString)]
        ),
        definition: knowledgeDefinition,
        result: knowledgeResult
      )
    )
    XCTAssertNil(
      AgentVerifiedAnswerRenderer.readAnswer(for: [mismatchedReceipt])
    )

    let reviewCall = AgentToolCall(
      name: "project_review",
      arguments: ["limit": .number(5)]
    )
    let reviewResult = await registry.execute(reviewCall, context: context)
    XCTAssertTrue(reviewResult.succeeded)
    XCTAssertLessThanOrEqual(reviewResult.modelText.count, 4_000)
    let reviewReceipt = try XCTUnwrap(
      AgentToolReceipt(
        runID: UUID(),
        call: try registry.normalize(reviewCall),
        definition: try XCTUnwrap(registry.definition(named: "project_review")),
        result: reviewResult
      )
    )
    XCTAssertTrue(
      try XCTUnwrap(
        AgentVerifiedAnswerRenderer.readAnswer(for: [reviewReceipt])
      ).contains("Project review")
    )

    let taskCall = AgentToolCall(
      name: "task_search",
      arguments: [
        "query": .string("release device"),
        "include_completed": .bool(false),
        "limit": .number(5),
      ]
    )
    let taskResult = await registry.execute(taskCall, context: context)
    let taskReceipt = try XCTUnwrap(
      AgentToolReceipt(
        runID: UUID(),
        call: try registry.normalize(taskCall),
        definition: try XCTUnwrap(registry.definition(named: "task_search")),
        result: taskResult
      )
    )
    XCTAssertTrue(
      try XCTUnwrap(
        AgentVerifiedAnswerRenderer.readAnswer(for: [taskReceipt])
      ).contains(task.id.uuidString)
    )

    let handoffCall = AgentToolCall(
      name: "handoff_get",
      arguments: ["id": .string(handoff.id.uuidString)]
    )
    let handoffResult = await registry.execute(handoffCall, context: context)
    let handoffReceipt = try XCTUnwrap(
      AgentToolReceipt(
        runID: UUID(),
        call: try registry.normalize(handoffCall),
        definition: try XCTUnwrap(registry.definition(named: "handoff_get")),
        result: handoffResult
      )
    )
    let handoffAnswer = try XCTUnwrap(
      AgentVerifiedAnswerRenderer.readAnswer(for: [handoffReceipt])
    )
    XCTAssertTrue(handoffAnswer.contains("Saved handoff"))
    XCTAssertTrue(handoffAnswer.contains("replayed no prior action"))
  }

  func testNewWritesRequireExplicitRoutesAndMutateOnlyTargetState()
    async throws
  {
    let now = Date(timeIntervalSince1970: 1_750_000_000)
    let task = AssistantTaskItem(
      title: "Recheck TestFlight",
      isCompleted: true,
      createdAt: now,
      completedAt: now
    )
    let handoff = SessionHandoff(
      title: "Release checkpoint",
      focus: "Continue release verification",
      summary: "No previous action should replay.",
      createdAt: now
    )
    var workspace = AssistantWorkspace.empty
    workspace.tasks = [task]
    workspace.handoffs = [handoff]
    workspace.activeHandoffID = handoff.id
    let store = ContinuityTestStore(workspace: workspace)
    let context = continuityExecutionContext(store: store)
    let registry = AgentToolRegistry()

    let reopenDefinition = try XCTUnwrap(
      registry.definition(named: "task_reopen")
    )
    XCTAssertEqual(
      AgentPolicy.decision(for: reopenDefinition, settings: AgentSettings()),
      .requireApproval(
        "This action changes saved assistant data on this iPhone."
      )
    )
    let reopened = await registry.execute(
      AgentToolCall(
        name: "task_reopen",
        arguments: ["id": .string(task.id.uuidString)]
      ),
      context: context
    )
    XCTAssertTrue(reopened.succeeded)
    var updated = await store.snapshot()
    XCTAssertFalse(try XCTUnwrap(updated.tasks.first).isCompleted)
    XCTAssertNil(updated.tasks.first?.completedAt)
    let reopenedAgain = await registry.execute(
      AgentToolCall(
        name: "task_reopen",
        arguments: ["id": .string(task.id.uuidString)]
      ),
      context: context
    )
    XCTAssertFalse(reopenedAgain.succeeded)

    let deactivated = await registry.execute(
      AgentToolCall(
        name: "handoff_deactivate",
        arguments: ["id": .string(handoff.id.uuidString)]
      ),
      context: context
    )
    XCTAssertTrue(deactivated.succeeded)
    updated = await store.snapshot()
    XCTAssertNil(updated.activeHandoffID)
    XCTAssertEqual(updated.handoffs, [handoff])
    let deactivatedAgain = await registry.execute(
      AgentToolCall(
        name: "handoff_deactivate",
        arguments: ["id": .string(handoff.id.uuidString)]
      ),
      context: context
    )
    XCTAssertFalse(deactivatedAgain.succeeded)
  }

  func testSchemaFiveRoundTripPreservesContinuityState() async throws {
    let fixedDate = Date(timeIntervalSince1970: 1_700_000_000)
    let item = KnowledgeItem(
      kind: .preference,
      title: "Verification style",
      content: "Keep compile evidence separate from device evidence.",
      tags: ["testing"],
      source: .user,
      createdAt: fixedDate
    )
    let handoff = KnowledgeEngine().makeHandoff(
      title: "Release handoff",
      focus: "Run the remaining device walkthrough.",
      knowledge: [item],
      tasks: [],
      createdAt: fixedDate
    )
    let digest = KnowledgeEngine().synthesizeContext(
      knowledge: [item],
      tasks: [],
      handoff: handoff,
      createdAt: fixedDate
    )
    var workspace = AssistantWorkspace.empty
    workspace.knowledgeItems = [item]
    workspace.knowledgeDigests = [digest]
    workspace.handoffs = [handoff]
    workspace.activeHandoffID = handoff.id

    let fileURL = FileManager.default.temporaryDirectory.appendingPathComponent(
      "dolphin-schema-five-\(UUID().uuidString).json"
    )
    defer { try? FileManager.default.removeItem(at: fileURL) }
    let persistence = try AssistantPersistence(fileURL: fileURL)
    try await persistence.save(workspace)
    let reloaded = try await persistence.load()
    XCTAssertEqual(reloaded, workspace)
  }
}

private actor ContinuityTestStore {
  private var workspace: AssistantWorkspace

  init(workspace: AssistantWorkspace) {
    self.workspace = workspace
  }

  func snapshot() -> AssistantWorkspace { workspace }

  func commit(
    _ mutation: AssistantWorkspaceMutation,
    resultBuilder: AssistantWorkspaceResultBuilder
  ) throws -> AgentToolResult {
    switch mutation {
    case .saveMemory(let item):
      workspace.memories.append(item)
      return try resultBuilder(.memory(item))
    case .captureKnowledge(let item):
      workspace.knowledgeItems.append(item)
      return try resultBuilder(.knowledge(item))
    case .createHandoff(let handoff):
      if let index = workspace.handoffs.firstIndex(where: { $0.id == handoff.id }) {
        workspace.handoffs[index] = handoff
      } else {
        workspace.handoffs.append(handoff)
      }
      workspace.activeHandoffID = handoff.id
      return try resultBuilder(.handoff(handoff))
    case .restoreHandoff(let id, let restoredAt):
      guard let index = workspace.handoffs.firstIndex(where: { $0.id == id }) else {
        throw ContinuityTestError.handoffNotFound
      }
      workspace.handoffs[index].restoredAt = restoredAt
      workspace.activeHandoffID = id
      return try resultBuilder(.handoff(workspace.handoffs[index]))
    case .deactivateHandoff(let id):
      guard workspace.activeHandoffID == id,
        let handoff = workspace.handoffs.first(where: { $0.id == id })
      else { throw ContinuityTestError.handoffNotFound }
      workspace.activeHandoffID = nil
      return try resultBuilder(.handoff(handoff))
    case .addTask(let item):
      workspace.tasks.append(item)
      return try resultBuilder(.task(item))
    case .completeTask(let id, let completedAt):
      guard let index = workspace.tasks.firstIndex(where: { $0.id == id }) else {
        throw ContinuityTestError.taskNotFound
      }
      workspace.tasks[index].isCompleted = true
      workspace.tasks[index].completedAt = completedAt
      return try resultBuilder(.task(workspace.tasks[index]))
    case .reopenTask(let id):
      guard let index = workspace.tasks.firstIndex(where: { $0.id == id }),
        workspace.tasks[index].isCompleted
      else { throw ContinuityTestError.taskNotFound }
      workspace.tasks[index].isCompleted = false
      workspace.tasks[index].completedAt = nil
      return try resultBuilder(.task(workspace.tasks[index]))
    }
  }
}

private enum ContinuityTestError: Error {
  case handoffNotFound
  case taskNotFound
}

private struct FixedContinuityDateProvider: DateProviding {
  let date: Date

  func now() -> Date { date }
}

private func continuityExecutionContext(
  store: ContinuityTestStore
) -> AgentToolExecutionContext {
  AgentToolExecutionContext(
    snapshot: { await store.snapshot() },
    commitLocalWrite: { mutation, _, _, resultBuilder in
      try await store.commit(mutation, resultBuilder: resultBuilder)
    }
  )
}